package top.fullj.security;

import com.google.common.io.BaseEncoding;
import top.fullj.io.FileStreamCompat;

import javax.crypto.Cipher;
import javax.crypto.EncryptedPrivateKeyInfo;
import javax.crypto.SecretKey;
import javax.crypto.SecretKeyFactory;
import javax.crypto.spec.PBEKeySpec;
import java.io.*;
import java.security.GeneralSecurityException;
import java.security.KeyException;
import java.security.KeyFactory;
import java.security.PrivateKey;
import java.security.cert.CertificateException;
import java.security.cert.CertificateFactory;
import java.security.cert.X509Certificate;
import java.security.spec.InvalidKeySpecException;
import java.security.spec.PKCS8EncodedKeySpec;
import java.util.ArrayList;
import java.util.List;
import java.util.regex.Matcher;
import java.util.regex.Pattern;

public abstract class PemReader {

    private static final Pattern CERT_PATTERN = Pattern.compile(
            "-+BEGIN\\s+.*CERTIFICATE[^-]*-+(?:\\s|\\r|\\n)+" + // Header
                    "([a-z0-9+/=\\r\\n]+)" +                    // Base64 text
                    "-+END\\s+.*CERTIFICATE[^-]*-+",            // Footer
            Pattern.CASE_INSENSITIVE);
    private static final Pattern KEY_PATTERN = Pattern.compile(
            "-+BEGIN\\s+.*PRIVATE\\s+KEY[^-]*-+(?:\\s|\\r|\\n)+" + // Header
                    "([a-z0-9+/=\\r\\n]+)" +                       // Base64 text
                    "-+END\\s+.*PRIVATE\\s+KEY[^-]*-+",            // Footer
            Pattern.CASE_INSENSITIVE);

    public static PrivateKey readPrivateKey(String keyFile, String keyPassword)
            throws IOException, GeneralSecurityException {
        try (InputStream is = FileStreamCompat.getInputStream(keyFile)) {
            return readPrivateKey(is, keyPassword);
        }
    }

    public static PrivateKey readPrivateKey(File keyFile, String keyPassword)
            throws IOException, GeneralSecurityException {
        try (InputStream is = new FileInputStream(keyFile)) {
            return toPrivateKey(is, keyPassword);
        } catch (FileNotFoundException e) {
            throw new KeyException("could not find key file: " + keyFile);
        }
    }

    public static PrivateKey readPrivateKey(InputStream keyStream, String keyPassword)
            throws IOException, GeneralSecurityException {
        return toPrivateKey(keyStream, keyPassword);
    }

    public static X509Certificate[] readX509Certificates(String file) throws CertificateException {
        try (InputStream is = FileStreamCompat.getInputStream(file)) {
            return readX509Certificates(is);
        } catch (FileNotFoundException e) {
            throw new CertificateException("could not find certificate file: " + file);
        } catch (IOException e) {
            // This is not expected to happen, but re-throw in case it does.
            throw new RuntimeException(e);
        }
    }

    public static X509Certificate[] readX509Certificates(File file) throws CertificateException {
        try (InputStream is = new FileInputStream(file)) {
            return toX509Certificates(is);
        } catch (FileNotFoundException e) {
            throw new CertificateException("could not find certificate file: " + file);
        } catch (IOException e) {
            // This is not expected to happen, but re-throw in case it does.
            throw new RuntimeException(e);
        }
    }

    public static X509Certificate[] readX509Certificates(InputStream is) throws CertificateException {
        return toX509Certificates(is);
    }

    private static X509Certificate[] toX509Certificates(InputStream stream) throws CertificateException {
        if (stream == null) {
            return null;
        }

        List<byte[]> certs = readCertificates(stream);

        CertificateFactory cf = CertificateFactory.getInstance("X.509");
        X509Certificate[] x509Certs = new X509Certificate[certs.size()];

        for (int i = 0; i < certs.size(); i++) {
            try (InputStream is = new ByteArrayInputStream(certs.get(i))) {
                x509Certs[i] = (X509Certificate) cf.generateCertificate(is);
            } catch (IOException e) {
                // This is not expected to happen, but re-throw in case it does.
                throw new RuntimeException(e);
            }
        }
        return x509Certs;
    }

    private static PrivateKey toPrivateKey(InputStream keyStream, String keyPassword)
            throws IOException, GeneralSecurityException {
        byte[] encodedKey = readPrivateKey(keyStream);

        PKCS8EncodedKeySpec encodedKeySpec = generateKeySpec(
                keyPassword == null ? null : keyPassword.toCharArray(), encodedKey);
        try {
            return KeyFactory.getInstance("RSA").generatePrivate(encodedKeySpec);
        } catch (InvalidKeySpecException ignore) {
            try {
                return KeyFactory.getInstance("DSA").generatePrivate(encodedKeySpec);
            } catch (InvalidKeySpecException ignore2) {
                try {
                    return KeyFactory.getInstance("EC").generatePrivate(encodedKeySpec);
                } catch (InvalidKeySpecException e) {
                    throw new InvalidKeySpecException("Neither RSA, DSA nor EC worked", e);
                }
            }
        }
    }

    protected static PKCS8EncodedKeySpec generateKeySpec(char[] password, byte[] key)
            throws IOException, GeneralSecurityException {
        if (password == null) {
            return new PKCS8EncodedKeySpec(key);
        }

        EncryptedPrivateKeyInfo encryptedPrivateKeyInfo = new EncryptedPrivateKeyInfo(key);
        SecretKeyFactory keyFactory = SecretKeyFactory.getInstance(encryptedPrivateKeyInfo.getAlgName());
        PBEKeySpec pbeKeySpec = new PBEKeySpec(password);
        SecretKey pbeKey = keyFactory.generateSecret(pbeKeySpec);

        Cipher cipher = Cipher.getInstance(encryptedPrivateKeyInfo.getAlgName());
        cipher.init(Cipher.DECRYPT_MODE, pbeKey, encryptedPrivateKeyInfo.getAlgParameters());

        return encryptedPrivateKeyInfo.getKeySpec(cipher);
    }

    private static byte[] readPrivateKey(InputStream is) throws KeyException {
        String content;
        try {
            content = readContent(is);
        } catch (IOException e) {
            throw new KeyException("failed to read key input stream", e);
        }

        Matcher m = KEY_PATTERN.matcher(content);
        if (!m.find()) {
            throw new KeyException("could not find a PKCS #8 private key in input stream" +
                    " (see https://netty.io/wiki/sslcontextbuilder-and-private-key.html for more information)");
        }
        String base64 = m.group(1).replaceAll("[\r\n]", "");
        return BaseEncoding.base64().decode(base64);
    }

    private static String readContent(InputStream is) throws IOException {
        try (ByteArrayOutputStream out = new ByteArrayOutputStream()) {
            byte[] buf = new byte[8192];
            for (;;) {
                int ret = is.read(buf);
                if (ret < 0) {
                    break;
                }
                out.write(buf, 0, ret);
            }
            return out.toString("US-ASCII");
        }
    }

    private static List<byte[]> readCertificates(InputStream is) throws CertificateException {
        String content;
        try {
            content = readContent(is);
        } catch (IOException e) {
            throw new CertificateException("failed to read certificate input stream", e);
        }

        List<byte[]> certs = new ArrayList<>();
        Matcher m = CERT_PATTERN.matcher(content);
        int start = 0;
        while (m.find(start)) {
            String base64 = m.group(1).replaceAll("[\r\n]", "");
            byte[] der = BaseEncoding.base64().decode(base64);
            certs.add(der);
            start = m.end();
        }

        if (certs.isEmpty()) {
            throw new CertificateException("found no certificates in input stream");
        }

        return certs;
    }


}
